'''
$LastChangedBy: csgillespie $
$LastChangedDate: 2007-06-06 11:08:06 +0000 (Wed, 06 Jun 2007) $
Colin Gillespie (c.gillespie@ncl.ac.uk)
Deals with the libsbml objects.
'''

import libsbml

#my packages
import ModifySBML
import DeleteSBML
import AddSBML
from GlobalVariables.GlobalAttributes import GlobalParameters

SBML_PARTS = ('Compartments', 'Species', 'Parameters', 'Reactions', 'Events')

class ModelHandler(GlobalParameters):
	'''This class privides a wrapper around some of the sbml functions
	
	In particular add/delete/modify SBML
	'''
	def __init__(self, base_values=None):
		GlobalParameters.__init__(self, base_values=base_values)
		self.m = self.d.getModel()
		
	def __str__(self):
		return self.d.toSBML()
	
	def __change_args_to_dict(self, **args):
		'''A private function that generates a dictionary used the add & modify 
		methods.
		It transforms the args: Compartments={key1,value1},Species={key2,value2}
		into a dictionary: {'Co':{key1,value1},'Sp':{key2,value2}}
		'''
		change_dict = {}
		for ele in SBML_PARTS:
			first_two_letters = ele[:2]
			change_dict[first_two_letters] = args.get(ele, {})
		return change_dict
	
	def __readInSbml(self, sbml, validation=None):
		'''Reads in an sbml string and returns
			a libsbml document
		'''
		if validation == 'BASIC':
			valid_level = libsbml.XML_SCHEMA_VALIDATION_BASIC
		else:
			valid_level = libsbml.XML_SCHEMA_VALIDATION_NONE
		
		"""Reads in the sbml and set validation parameters"""
		read = libsbml.SBMLReader()
		read.setSchemaFilenameL1v1(self.path + 'xml_schemas/sbml/sbml-l1v1.xsd')
		read.setSchemaFilenameL1v2(self.path + 'xml_schemas/sbml/sbml-l1v2.xsd')
		read.setSchemaFilenameL2v1(self.path + 'xml_schemas/sbml/sbml-l2v1.xsd')
		read.setSchemaValidationLevel(valid_level)
		
		return read.readSBMLFromString(sbml)
		
		
	def isValid(self):
		'''Returns 1 is the model is valid. 
			Returns 0 for an invalid model.
		'''
		if not self.path:
			print 'Unable to locate the style-sheets. Change the PATH'
			print 'variable in pysbml.py'
			return None
			
		sbml_validated = self.__readInSbml(self.d.toSBML(), 'BASIC')
		
		num_warnings = sbml_validated.getNumWarnings()
		num_errors = sbml_validated.getNumErrors()
		num_fatals = sbml_validated.getNumFatals()
		if (num_warnings + num_errors + num_fatals) == 0:
			return 1
		else:
			return 0
		
	def validate(self):
		'''Validates the model. 
		If invalid then returns an output of the errors. 
		If valid returns a summary
		'''
		if not self.path:
			print 'Unable to locate the style-sheets. Change the PATH'
			print 'variable in pysbml.py'
			return None
		
		d_validated = self.__readInSbml(self.d.toSBML(), 'BASIC')
		
		num_warnings = d_validated.getNumWarnings()
		num_errors = d_validated.getNumErrors()
		num_fatals = d_validated.getNumFatals()			
		out = ""
		if (num_warnings + num_errors + num_fatals) == 0:
			out += "Valid SBML model\n"
			model = d_validated.getModel()
			out += "".join("\nModel summary\n\n")
			out += "".join("%s\n" % model.getId())
			out += "".join("Compartments %d \n" % model.getNumCompartments())
			out += "".join("Unit Definitions %d \n" % model.getNumUnitDefinitions())
			out += "".join("Species %d \n" % model.getNumSpecies())
			out += "".join("Reactions %d \n" % model.getNumReactions())
			out += "".join("Rules %d \n" % model.getNumRules())
			out += "".join("Parameters %d \n" % model.getNumParameters())
			out += "".join("Events %d \n" % model.getNumEvents())
		else:
			out += "Invalid SBML model\n"
			out = "%s Warnings\n" % num_warnings
			for warning_num in range(num_warnings):
				warn = d_validated.getWarning(warning_num)
				out += "".join("line: %d column: %d  warning: %s"
					% (warn.getLine(), warn.getColumn(), warn.getMessage()))
		
			out +="%s Errors\n" % num_errors
			for error_num in range(num_errors):
				error = d_validated.getError(error_num)
				out += "".join("line: %d Column: %d  Error: %s\n" 
					% (error.getLine(), error.getColumn(), error.getMessage()))
		
			out += "%s Fatals\n" % num_fatals
			for fatal_num in range(num_fatals):
				fatal = d_validated.getFatal(fatal_num)
				out += "".join("line: %d Column: %d  Fatal: %s" %
					 (fatal.getLine(), fatal.getColumn(), fatal.getMessage()))
		return out
	
	def convertor(self, level, version):
		'''Convert begin sbml versions'''
		
		if level not in (1, 2):
			return "invalid level"
		if version not in (1, 2):
			return "invalid version"
		if self.isValid() != 0:
			return self.isValid()
		
		if self.d.getLevel() != level:
			self.d.setLevel(level)
			self.d.setVersion(version)
			self.sbml = self.d.toSBML()
		elif self.d.getVersion() != version:
			self.d.setVersion(version)
			self.sbml = self.d.toSBML()

		return self.d.toSBML()
	
	def add(self, **args):
		'''Add to the current SBML model
		Usage:
			add(Compartments={'Id':'Cell1', 'Size':10},
			Species={'Id':'ATP1','Compartment':'Cell1'},
			Parameters={'Id':'ZZZZZ','Value':10},
			Reactions = {'Id':'NewReaction',
				'listOfReactants':[{'Species':'ATP'}, {'Species':'ADP'}], 
				'listOfProducts':[{'Species':'ATP'}, {'Species':'ADP'}],
				'kineticLaw':{'Formula':2}},
			Events = {'Id':'NewEvent','Trigger':'eq(ATP,ADP)',
				'listOfEventAssignments':[{'Variable':'T', 'Formula':'2*T'}]}
			)
		
		This would add in a new Cell with Id 'Cell1' and Size 10
		'''
		add_dict = self.__change_args_to_dict(**args)
		AddSBML.AddSBML(self.m, SBML_PARTS, add_dict).alterModel()
		return 1	
	
	def modify(self, **args):
		'''Modifies the sbml model
		
		Usage:
			modify(Compartments={'Cell':{'Id':'newCell'}}
		This would change the id of compartment cell to newCell
				
		'''
		recursive = args.get('recursive', 1)
		alter_dict = self.__change_args_to_dict(**args)
		
		for mod_id, mod_new_id in args.get('Ids', {}).items():
			first_two_letters = self.__searchForIds(mod_id)[:2]
			alter_dict[first_two_letters][mod_id] = {'Id':mod_new_id}

		ModifySBML.ModifySBML(self.m, SBML_PARTS, alter_dict, \
			recursive = recursive).alterModel()
		return 1	
	
	def delete(self, **args):
		'''Deletes part of the sbml model
		Usage:
			delete(Species=['ATP'],Reaction=['R1','R2']
		
		This would delete species ATP and reactions R1 and R2
		Also any reaction or event that depended on ATP would
		be deleted
		
		if recursive==0 then only a single element is delete.
		and NOT the dependencies.
		'''
		recursive = args.get('recursive', 1)
		alter_dict = {}
		for ele in SBML_PARTS:
			first_two_letters = ele[:2]
			alter_dict[first_two_letters] = args.get(ele, [])
		
		
		for del_id in args.get('Ids', []):
			first_two_letters = self.__searchForIds(del_id)[:2]
			alter_dict[first_two_letters].append(del_id)
		
		DeleteSBML.DeleteSBML(self.m, SBML_PARTS, alter_dict, \
			recursive = recursive).alterModel()
		return 1
	
	def __searchForIds(self, del_id):
		'''Used in the delete and modify methods.
		If given an SBML id, we need to find out what type of element it is.
		e.g. Is it species, model, etc.
		'''
		for ele in SBML_PARTS:
			for sbml_id in getattr(self.m, 'getListOf'+ele)():
				if sbml_id.getId() == del_id:
					return ele
		raise 'Invalid Id'
	
	def toSBML(self):
		'''Prints the sbml'''
		return self.d.toSBML()
	
	def setId(self, sbml_id):
		'''Sets the sbml id'''
		self.m.setId(sbml_id)
	
	def setNotes(self, notes):
		'''set sbml notes'''
		self.m.setNotes(notes)
	
	#All these methods below should be able to be automatically generated
	#Unfortunately, I couldn't do it :(
	
	def addCompartments(self, **kwds):
		add_dict = self.__change_args_to_dict(Compartments=dict(**kwds))
		AddSBML.AddSBML(self.m, SBML_PARTS, add_dict).alterModel()
		return 1
				
	def addSpecies(self, **kwds):
		add_dict = self.__change_args_to_dict(Species=dict(**kwds))
		AddSBML.AddSBML(self.m, SBML_PARTS, add_dict).alterModel()
		return 1
	
	def addParameters(self, **kwds):
		add_dict = self.__change_args_to_dict(Parameters=dict(**kwds))
		AddSBML.AddSBML(self.m, SBML_PARTS, add_dict).alterModel()
		return 1
		
	def addReactions(self, **kwds):
		add_dict = self.__change_args_to_dict(Reactions=dict(**kwds))
		AddSBML.AddSBML(self.m, SBML_PARTS, add_dict).alterModel()
		return 1
		
	def addEvents(self, **kwds):
		add_dict = self.__change_args_to_dict(Events=dict(**kwds))
		AddSBML.AddSBML(self.m, SBML_PARTS, add_dict).alterModel()
		return 1
	
	#Modify methods
	
	def modifyCompartments(self, **kwds):
		recursive = kwds.get('recursive', 1)
		alter_dict = self.__change_args_to_dict(Compartments=dict(**kwds))
		for mod_id, mod_new_id in kwds.get('Ids', {}).items():
			first_two_letters = self.__searchForIds(mod_id)[:2]
			alter_dict[first_two_letters][mod_id] = {'Id':mod_new_id}
	
		ModifySBML.ModifySBML(self.m, SBML_PARTS, alter_dict, \
			recursive = recursive).alterModel()
		return 1
		
	def modifySpecies(self, **kwds):
		recursive = kwds.get('recursive', 1)
		alter_dict = self.__change_args_to_dict(Species=dict(**kwds))
		for mod_id, mod_new_id in kwds.get('Ids', {}).items():
			first_two_letters = self.__searchForIds(mod_id)[:2]
			alter_dict[first_two_letters][mod_id] = {'Id':mod_new_id}
	   
		ModifySBML.ModifySBML(self.m, SBML_PARTS, alter_dict, \
			recursive = recursive).alterModel()
		return 1
		
	def modifyParameters(self, **kwds):
		recursive = kwds.get('recursive', 1)
		alter_dict = self.__change_args_to_dict(Parameters=dict(**kwds))
		for mod_id, mod_new_id in kwds.get('Ids', {}).items():
			first_two_letters = self.__searchForIds(mod_id)[:2]
			alter_dict[first_two_letters][mod_id] = {'Id':mod_new_id}
	   
		ModifySBML.ModifySBML(self.m, SBML_PARTS, alter_dict, \
			recursive = recursive).alterModel()
		return 1
		
	def modifyReactions(self, **kwds):
		recursive = kwds.get('recursive', 1)
		alter_dict = self.__change_args_to_dict(Reactions=dict(**kwds))
		for mod_id, mod_new_id in kwds.get('Ids', {}).items():
			first_two_letters = self.__searchForIds(mod_id)[:2]
			alter_dict[first_two_letters][mod_id] = {'Id':mod_new_id}
	   
		ModifySBML.ModifySBML(self.m, SBML_PARTS, alter_dict, \
			recursive = recursive).alterModel()
		return 1
		
	def modifyEvents(self, **kwds):
		recursive = kwds.get('recursive', 1)
		alter_dict = self.__change_args_to_dict(Events=dict(**kwds))
		for mod_id, mod_new_id in kwds.get('Ids', {}).items():
			first_two_letters = self.__searchForIds(mod_id)[:2]
			alter_dict[first_two_letters][mod_id] = {'Id':mod_new_id}
	   
		ModifySBML.ModifySBML(self.m, SBML_PARTS, alter_dict, \
			recursive = recursive).alterModel()
		return 1
	
	
	#Delete Methods
	def __generate_empty_dictionary_for_delete_method(self):
		alter_dict = {}
		for ele in SBML_PARTS:
			first_two_letters = ele[:2]
			alter_dict[first_two_letters] = []
		
		return alter_dict
		
	def deleteCompartments(self, listOf, recursive=1):
		alter_dict = self.__generate_empty_dictionary_for_delete_method()
		alter_dict['Co'] = listOf
		
		DeleteSBML.DeleteSBML(self.m, SBML_PARTS, alter_dict, \
			recursive = recursive).alterModel()
		return 1
	
	def deleteSpecies(self, listOf, recursive=1):
		alter_dict = self.__generate_empty_dictionary_for_delete_method()
		alter_dict['Sp'] = listOf
		
		DeleteSBML.DeleteSBML(self.m, SBML_PARTS, alter_dict, \
			recursive = recursive).alterModel()
		return 1
	
	def deleteParameters(self, listOf, recursive=1):
		alter_dict = self.__generate_empty_dictionary_for_delete_method()
		alter_dict['Pa'] = listOf
		
		DeleteSBML.DeleteSBML(self.m, SBML_PARTS, alter_dict, \
			recursive = recursive).alterModel()
		return 1
	
	def deleteReactions(self, listOf, recursive=1):
		alter_dict = self.__generate_empty_dictionary_for_delete_method()
		alter_dict['Re'] = listOf
		
		DeleteSBML.DeleteSBML(self.m, SBML_PARTS, alter_dict, \
			recursive = recursive).alterModel()
		return 1
		
	def deleteEvents(self, listOf, recursive=1):
		alter_dict = self.__generate_empty_dictionary_for_delete_method()
		alter_dict['Ev'] = listOf
		
		DeleteSBML.DeleteSBML(self.m, SBML_PARTS, alter_dict, \
			recursive = recursive).alterModel()
		return 1
